import sqlite3
import re
import logging
from typing import List, Tuple, Any, Dict
from src.base.agent import Agent
from tabulate import tabulate

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler('text2sql.log'), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

class Text2SQLConverter:
    def __init__(self):
        """初始化Text2SQL转换器"""
        self.agent = Agent(
            name="SQLAgent",
            instructions="""你是一个专业的数据库查询专家，负责将自然语言转换为SQL查询语句。

数据库架构：
1. papers表：
   - id: INTEGER PRIMARY KEY
   - title: TEXT (论文标题)
   - authors: TEXT (作者)
   - conference: TEXT (发表会议)
   - year: INTEGER (发表年份)
   - citations: INTEGER (引用次数)
   - abstract: TEXT (摘要)

2. keywords表：
   - id: INTEGER PRIMARY KEY
   - paper_id: INTEGER (外键关联papers表)
   - keyword: TEXT (关键词)

请遵循以下规则：
1. 只返回SQL查询语句，不要包含其他解释
2. 使用标准SQL语法
3. 支持复杂查询，包括：
   - 多表连接
   - 条件过滤
   - 排序和分组
   - 聚合函数
4. 确保查询性能优化
"""
        )
        self.conn = sqlite3.connect(':memory:')
        self.cursor = self.conn.cursor()
        self._init_database()

    def _init_database(self):
        """初始化数据库表结构和示例数据"""
        # 创建论文表
        self.cursor.execute('''
        CREATE TABLE papers (
            id INTEGER PRIMARY KEY,
            title TEXT,
            authors TEXT,
            conference TEXT,
            year INTEGER,
            citations INTEGER,
            abstract TEXT
        )''')

        # 创建关键词表
        self.cursor.execute('''
        CREATE TABLE keywords (
            id INTEGER PRIMARY KEY,
            paper_id INTEGER,
            keyword TEXT,
            FOREIGN KEY (paper_id) REFERENCES papers(id)
        )''')

        # 插入示例论文数据
        papers = [
            (1, 'Attention is All You Need', 'Vaswani et al.', 'NeurIPS', 2017, 1000, 
             'We propose a new network architecture, the Transformer...'),
            (2, 'BERT: Pre-training of Deep Bidirectional Transformers', 'Devlin et al.', 'NAACL', 2019, 800,
             'We introduce a new language representation model called BERT...'),
            # ... 更多示例数据 ...
        ]
        self.cursor.executemany('INSERT INTO papers VALUES (?,?,?,?,?,?,?)', papers)

        # 插入关键词数据
        keywords = [
            (1, 1, 'attention'),
            (2, 1, 'transformer'),
            (3, 2, 'bert'),
            (4, 2, 'pre-training'),
            # ... 更多关键词 ...
        ]
        self.cursor.executemany('INSERT INTO keywords VALUES (?,?,?)', keywords)
        self.conn.commit()

    def clean_sql_query(self, sql_query: str) -> str:
        """清理SQL查询语句"""
        return re.sub(r'```sql\s*|\s*```', '', sql_query).strip()

    def execute_sql(self, sql_query: str) -> Tuple[List[Any], str]:
        """执行SQL查询并返回结果和可能的错误信息"""
        try:
            self.cursor.execute(sql_query)
            results = self.cursor.fetchall()
            return results, None
        except sqlite3.Error as e:
            error_msg = f"SQL执行错误: {str(e)}"
            logger.error(error_msg)
            return [], error_msg

    def format_results(self, results: List[Any], error: str = None) -> str:
        """格式化查询结果"""
        if error:
            return f"错误: {error}"
            
        if not results:
            return "未找到匹配的结果。"

        headers = [description[0] for description in self.cursor.description]
        return tabulate(results, headers=headers, tablefmt="grid")

    def process_query(self, natural_language_query: str) -> Dict[str, str]:
        """处理自然语言查询并返回结果"""
        try:
            # 生成SQL查询
            response = self.agent.run(
                messages=[{"role": "user", "content": natural_language_query}]
            )
            sql_query = self.clean_sql_query(response["result"])
            logger.info(f"生成的SQL查询: {sql_query}")

            # 执行查询
            results, error = self.execute_sql(sql_query)
            
            # 格式化结果
            formatted_results = self.format_results(results, error)

            return {
                "query": natural_language_query,
                "sql": sql_query,
                "results": formatted_results,
                "error": error
            }
        except Exception as e:
            error_msg = f"处理查询时出错: {str(e)}"
            logger.error(error_msg)
            return {
                "query": natural_language_query,
                "sql": "",
                "results": "",
                "error": error_msg
            }

    def close(self):
        """关闭数据库连接"""
        self.conn.close()

def main():
    """主程序入口"""
    converter = Text2SQLConverter()
    logger.info("Text2SQL系统已启动")
    
    print("欢迎使用Text2SQL转换系统")
    print("输入 'exit' 或 'quit' 退出程序")
    print("示例查询：")
    print("1. 查找2018年之后发表的论文")
    print("2. 统计每个会议的论文数量")
    print("3. 找出引用次数最多的5篇论文")
    
    while True:
        user_input = input("\n请输入您的查询 (或 'exit' 退出): ")
        if user_input.lower() in ['exit', 'quit']:
            break
            
        result = converter.process_query(user_input)
        print("\nSQL查询:")
        print(result["sql"])
        print("\n查询结果:")
        print(result["results"])
        
        if result["error"]:
            print("\n错误信息:")
            print(result["error"])
    
    converter.close()
    logger.info("Text2SQL系统已关闭")

if __name__ == "__main__":
    main()